# -*- coding: utf-8 -*-

import copy
import time
from abc import ABC

import pandas as pd

from entity.entity import Entity
from entity.constant import *
from ss.istrategy import IStrategy
from entity.stra_entity import Tick, Bar


class DFBar(pd.DataFrame, ABC):
    """pd.DataFrame格式的K线序列
    将索引变为int， 原 datetime 类索引变为一列
    """

    def __init__(self, data=None, columns=None, index=None):
        super(DFBar, self).__init__(data, index=index)
        if columns is None:
            columns = {'Open': 'open', 'High': 'high', 'Low': 'low', 'Close': 'close', 'Volume': 'volume',
                       'Amount': 'amount', 'Interest': 'interest', 'DateTime': 'datetime', 'OpenInterest': 'oi',
                       'Datetime': 'datetime', 'date': 'datetime', 'Date': 'datetime'}
        self.rename(columns=columns, inplace=True)
        if data is None:
            return
        if self.index.name not in ['datetime', 'DateTime', 'Datetime']:
            self.set_index('datetime', inplace=True, drop=False)
        elif 'datetime' not in self.columns:
            self.index.name = 'datetime'
            self.reset_index(inplace=True, drop=False)
            self.set_index('datetime', inplace=True, drop=False)

    def add_data(self, data, index=None):
        if isinstance(data, DFBar):
            return self.append(data)
        else:
            df = DFBar(data=data, index=index)
            return self.append(df)


class Bars(Entity):
    """pd.DataFrame格式的K线序列
    将索引变为int， 原 datetime 类索引变为一列
    """

    def __init__(self, data=None, index_name='datetime',
                 columns=None, period=60, interval_type=IntervalType.Minute, parent=None, instrument_id=''):
        self.parent = parent  # instrument.manager.StraInstr
        self.data = DFBar(data, columns=columns)
        self.columns = columns
        self.index_name = index_name

        self.instrument_id = instrument_id  # if parent is None else parent.id
        self.period = period
        self.interval_type = interval_type  # bar 所属周期,按分钟计：0==Tick, 60==1分钟，60*3==3分钟，60*60==1小时，60*60*24==日
        # 列简称
        self.simple_sgn = ['d', 'o', 'h', 'l', 'c', 'v', 'a', 'i', 'D', 'O', 'H', 'L', 'C', 'V', 'A', 'I']

        '''Strategies: key=stra.id, value=stra_instance'''
        self.stras = {}
        '''Indicators'''
        self.indis = {}
        # current tick
        self.tick: Tick = Tick()
        # current bar
        self.bar: Bar = None

    def add_data(self, data, index=None):
        if self.data.empty:
            self.data = DFBar(data, index=index)
        else:
            self.data.add_data(data, index=index)

        for stra in self.stras.values():
            stra.on_bars(self)

    def add_strategy(self, strategy: IStrategy):
        '''add strategy '''
        strategy.set_instrument_id = self.instrument_id
        self.stras[strategy.id] = strategy
        if self.data.empty:
            return
        strategy.on_bars(self)

    def add_indicator(self, indicator):
        ''' add indicators '''
        self.indis[indicator.name] = indicator

    @property
    def length(self):
        return len(self.data)

    def ref(self, *args):
        c_name = 'close'
        ref_num = 1
        if len(args) == 2:
            c_name = args[0]
            ref_num = args[1]
        elif len(args) == 1:
            if isinstance(args[0], int):
                c_name = 'close'
                ref_num = args[0]
            elif isinstance(args[0], str):
                c_name = args[0]
                ref_num = 1

        return self.data[c_name].iloc[-ref_num]

    def refo(self, num=1):
        return self.data['open'].iloc[-num]

    def refh(self, num=1):
        return self.data['high'].iloc[-num]

    def refl(self, num=1):
        return self.data['low'].iloc[-num]

    def refc(self, num=1):
        return self.data['close'].iloc[-num]

    def refv(self, num=1):
        return self.data['volume'].iloc[-num]

    def refa(self, num=1):
        return self.data['amount'].iloc[-num]

    def refi(self, num=1):
        return self.data['interest'].iloc[-num]

    def refd(self, num=1):
        return self.data['datetime'].iloc[-num]

    @property
    def d(self):
        return self.data['datetime']

    @property
    def o(self):
        return self.data['open']

    @property
    def h(self):
        return self.data['high']

    @property
    def l(self):
        return self.data['low']

    @property
    def c(self):
        return self.data['close']

    @property
    def v(self):
        return self.data['volume']

    @property
    def a(self):
        return self.data['amount']

    @property
    def i(self):
        return self.data['interest']

    @property
    def oi(self):
        return self.data['oi']

    D = d
    O = o
    H = h
    L = l
    C = c
    V = v
    A = a
    I = i
    Close = c
    High = h
    Low = l
    Open = o
    Volume = vol = v

    def last(self, num=1):
        """最后的Bars"""
        return Bar(bar=self.data.iloc[-num].to_dict())

    def first(self, num=1):
        """最前边的 Bars """
        return Bar(self.data.iloc[num - 1].to_dict())

    def on_bar_update(self, bar):
        for ind in self.indis.values():
            ind.on_bar(bar)
        for stra in self.stras.values():
            stra.on_bar(bar)

    def on_bar(self, event):
        # 根据新的Bar数据，增加新的Bar
        bar = event.dict_['data']
        self.data[bar.datetime] = bar
        self.on_bar_update(bar)

    def on_tick_update(self, tick):
        for ind in self.indis:
            pass
            # ind.on_tq_tick(tick)
        for stra in self.stras:
            pass
            # stra.on_tq_tick(tick)

    def on_tick(self, event):
        # 收到新的tick数据时执行,根据新的Tick数据，更新Bar值
        tick = event.dict_['data']

        self.on_tick_update(tick)
        self.deal_tick(tick)

    def deal_tick(self, tick):
        '''分笔数据处理'''
        # 避免相同tick重复调用
        if self.tick.UpdateTime == tick.UpdateTime and self.tick.Volume == tick.Volume:
            return
        self.tick = Tick(tick.__dict__)
        tradingday = tick.TradingDay
        # 取此tick对应的分钟时间, bar的datetime是第一个数据的时间
        # bar_time = time.strptime(time.strftime("%Y-%m-%d %H:%M", tick.UpdateTime), "%Y-%m-%d %H:%M")
        bar_time = tick.UpdateTime[
                   :-2] + '00'  # time.strftime("%Y%m%d %H:%M:00", time.strptime(tick.UpdateTime, "%Y%m%d %H:%M:%S"))
        bar_datetime = f'{tick.ActionDay} {bar_time}'

        if self.length == 0 or self.bar is None:  # 新数据
            # bar_time, ins, h, l, o, c, v, i, a)
            bar = Bar(bar_datetime, tick.InstrumentID, tick.LastPrice, tick.LastPrice, tick.LastPrice, tick.LastPrice,
                      tick.Volume, tick.OpenInterest, 0, tradingday)
            if self.length > 0:
                if self.bar is None or self.bar.pre_v == 0:  # 实时行情首K即为新的分钟
                    bar.volume = 0
                else:
                    bar.volume = tick.Volume - self.bar.pre_v - self.bar.V
            bar._pre_volume = tick.Volume

            self.__new_bar__(bar)  # 新K线数据插入
        elif self.bar.D != bar_datetime:  # 新数据
            # bar_time, ins, h, l, o, c, v, i, a)
            bar = Bar(bar_datetime, tick.InstrumentID, tick.LastPrice, tick.LastPrice, tick.LastPrice, tick.LastPrice,
                      tick.Volume, tick.OpenInterest, 0, tradingday)
            if self.bar.pre_v == 0:  # 实时行情首K即为新的分钟
                bar.volume = 0
            else:
                bar.volume = tick.Volume - self.bar.pre_v - self.bar.V
            bar._pre_volume = tick.Volume

            self.__new_bar__(bar)  # 新K线数据插入
        else:  # bar 走完
            bar = self.bar
            bar.high = max(bar.H, tick.LastPrice)
            bar.low = min(bar.L, tick.LastPrice)
            bar.close = tick.LastPrice
            # 当时从服务器取到的数据,与ctp实时数据处于同一分钟,需做衔接处理.
            if bar.pre_v == 0:
                bar.pre_v = tick.Volume - bar.V  # 此tick产生的成交量忽略
            else:
                bar.volume = tick.Volume - bar.pre_v
            # bar._pre_volume = tick.Volume
            bar.interest = tick.OpenInterest
            bar.ap = tick.AveragePrice

    def __new_bar__(self, bar2):
        """有新min_bar添加"""
        bar = copy.copy(bar2)
        # print(bar.__dict__)
        bar_time = time.strptime(bar.D, "%Y%m%d %H:%M:%S")
        year = bar_time.tm_year
        mon = bar_time.tm_mon
        day = bar_time.tm_mday
        hour = bar_time.tm_hour
        mins = bar_time.tm_min
        # 由于 period 和 IntervalType.Minute 是以秒计的，要将 period 换算成分钟：
        period_minute = self.period // IntervalType.Minute
        if self.interval_type == IntervalType.Minute:
            mins = bar_time.tm_min // period_minute * period_minute  # 确保是 period_minute 的整数倍
        elif self.interval_type == IntervalType.Hour:
            hour = hour * IntervalType.Hour // self.period
            mins = 0
        elif self.interval_type == IntervalType.Day:
            hour = 0
            mins = 0
        elif self.interval_type == IntervalType.Month:
            hour = 0
            mins = 0
            day = 1
        elif self.interval_type == IntervalType.Year:
            hour = 0
            mins = 0
            day = 1
            mon = 1
        elif self.interval_type == IntervalType.Week:
            hour = 0
            mins = 0
            # 用周号替换日期
            day = time.strftime('%W', bar_time)

        # time -> str
        bar_time = f'{year}{mon:02d}{day:02d} {hour:02d}:{mins:02d}:00'
        if self.length == 0 or self.bar is None:
            bar.datetime = bar_time
            self.bar = bar
        elif self.bar.D != bar_time:
            # 将产生新BAR，将已经生成完整的保存，发送，更换为新的
            self.data = self.data.add_data(self.bar.__dict__, [self.bar.datetime])
            self.on_bar_update(self.bar)
            bar.datetime = bar_time
            self.bar = bar
        else:
            old_bar = self.bar
            old_bar.high = max(bar.H, old_bar.H)
            old_bar.low = min(bar.L, old_bar.L)
            old_bar.close = bar.C
            old_bar.volume += bar.V
            # self.V[-1] = old_bar.V
            old_bar.interest = bar.I
            # bar.A = tick.AveragePrice

    def ma(self, sgn='close', pma=20):
        ma_name = ''.join(['ma', sgn, str(pma)])
        ma_ = self.indis.get(ma_name, None)
        if not ma_:
            ma_ = self.data[sgn].rolling(window=pma, min_periods=1).mean()

        self.indis[ma_name] = {'ma': {'sgn': sgn, 'pma': pma}}
        self[ma_name] = ma_
        return ma_

    def add_indi(self, indi):
        self.indis[indi.id] = indi

    def add_stra(self, stra):
        self.stras[stra.id] = stra


class KLine(object):
    """"""

    def __init__(self, data=None, index_name='datetime',
                 columns=None, period=0, bars=None, symbol=''):
        if bars:
            self.bars = bars
        else:
            self.bars = Bars(data, index_name, columns, period)
        self.symbol = symbol
        self.period = period
        self.stras = self.bars.stras
        self.indis = self.bars.indis

    def add_indicator(self, indicator):
        self.bars.add_indicator(indicator)

    def add_strategy(self, strategy):
        self.bars.add_strategy(strategy)

    def plot(self):
        """"""


if __name__ == '__main__':
    # from data.get_data_hf_zmq import get_frame_bar
    from datetime import datetime, timedelta

    ed = datetime.now().strftime('%Y%m%d')
    sd = (datetime.now() - timedelta(days=10)).strftime('%Y%m%d')
    iid = 'v2001'

    # data = get_frame_bar(iid, sd, ed)
    # print(data.columns)
    # print(data.head())
    # fb = Bars(data=data)
    # print(type(data), len(data))
    # print(fb.data.index.name)
    # print(fb.data.head())
    # print(fb.data.columns)
    # print(fb.d[-1], fb.refd(1))
    # print(fb.h, '\n', type(fb.h))
    # print(fb.refh(2))

    # fb.ma()
    # print('maclose:', fb.maclose20)
